"""
WebSocket API Endpoints
Handles WebSocket connections for real-time communication during PPT generation.
"""

import logging
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import Optional
import asyncio
import json

from ..services.websocket_manager import connection_manager

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/ws/{session_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    session_id: str,
    token: Optional[str] = Query(None)
):
    """
    WebSocket endpoint for real-time communication during PPT generation
    
    Args:
        websocket: WebSocket connection instance
        session_id: Unique session identifier for the PPT generation
        token: Optional authentication token (for future use)
    """
    logger.info(f"WebSocket connection attempt for session: {session_id}")
    
    # Connect to the WebSocket
    connection_success = await connection_manager.connect(websocket, session_id)
    
    if not connection_success:
        logger.error(f"Failed to establish WebSocket connection for session: {session_id}")
        return
    
    try:
        # Keep the connection alive and handle incoming messages
        while True:
            try:
                # Wait for messages from client
                data = await websocket.receive_text()
                
                # Parse the message
                try:
                    message = json.loads(data)
                    await handle_client_message(session_id, message)
                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON received from session {session_id}: {data}")
                    await connection_manager.send_error_message(
                        session_id, 
                        "Invalid message format", 
                        "INVALID_JSON"
                    )
                
            except WebSocketDisconnect:
                logger.info(f"WebSocket disconnected for session: {session_id}")
                break
            except Exception as e:
                logger.error(f"Error handling WebSocket message for session {session_id}: {e}")
                await connection_manager.send_error_message(
                    session_id,
                    "Internal server error",
                    "INTERNAL_ERROR"
                )
                break
                
    except Exception as e:
        logger.error(f"WebSocket connection error for session {session_id}: {e}")
    finally:
        # Clean up the connection
        await connection_manager.disconnect(session_id)


async def handle_client_message(session_id: str, message: dict):
    """
    Handle incoming messages from WebSocket clients
    
    Args:
        session_id: Session identifier
        message: Parsed message from client
    """
    message_type = message.get("type")
    
    if message_type == "ping":
        # Respond to ping with pong
        await connection_manager.send_personal_message(session_id, {
            "type": "pong",
            "original_timestamp": message.get("timestamp")
        })
        
    elif message_type == "get_status":
        # Client requesting current status
        # This would typically query the generation session status
        await connection_manager.send_personal_message(session_id, {
            "type": "status_response",
            "status": "connected",
            "session_id": session_id
        })
        
    elif message_type == "cancel_generation":
        # Client requesting to cancel generation
        logger.info(f"Generation cancellation requested for session: {session_id}")
        await connection_manager.send_personal_message(session_id, {
            "type": "cancellation_acknowledged",
            "message": "Generation cancellation request received"
        })
        # TODO: Implement actual generation cancellation logic
        
    else:
        logger.warning(f"Unknown message type '{message_type}' from session {session_id}")
        await connection_manager.send_error_message(
            session_id,
            f"Unknown message type: {message_type}",
            "UNKNOWN_MESSAGE_TYPE"
        )


@router.get("/ws/status")
async def websocket_status():
    """
    Get WebSocket connection status and statistics
    
    Returns:
        dict: Connection statistics and status information
    """
    return {
        "active_connections": connection_manager.get_connection_count(),
        "status": "operational",
        "message": "WebSocket service is running"
    }


@router.post("/ws/broadcast")
async def broadcast_message(message: dict):
    """
    Broadcast a message to all active WebSocket connections
    (Admin endpoint - would typically require authentication)
    
    Args:
        message: Message to broadcast
        
    Returns:
        dict: Broadcast result
    """
    await connection_manager.broadcast_message(message)
    return {
        "status": "success",
        "message": "Message broadcasted to all active connections",
        "connections_count": connection_manager.get_connection_count()
    }


@router.get("/ws/connections/{session_id}")
async def get_connection_info(session_id: str):
    """
    Get connection information for a specific session
    
    Args:
        session_id: Session identifier
        
    Returns:
        dict: Connection information or error
    """
    if not connection_manager.is_connected(session_id):
        return {
            "status": "not_connected",
            "session_id": session_id,
            "message": "No active connection for this session"
        }
    
    connection_info = connection_manager.get_connection_info(session_id)
    return {
        "status": "connected",
        "session_id": session_id,
        "connection_info": connection_info
    }